#!/usr/bin/env python
# encoding: utf-8
'''
@author: wangjianrong
@software: pycharm
@file: lenet.py
@time: 2020/10/15 11:54
@desc:
'''
import torch
from torch_op.common import init_weights


class Lenet(torch.nn.Module):
    def __init__(self):
        super(Lenet,self).__init__()
        self.conv1 = torch.nn.Conv2d(1,6,3,1)
        self.conv2 = torch.nn.Conv2d(6,16,3,1)
        self.fc1 = torch.nn.Linear(9216,120)
        self.fc2 = torch.nn.Linear(120,84)
        self.fc3 = torch.nn.Linear(84,10)
        init_weights(self)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0),-1)
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x
        # return self.fc3(self.fc2(self.fc1(self.conv2(self.conv1(x)))))